import numpy as np
import cv2
import os
from getStrokeLabel import genStrokes
from config import Config
config = Config()
draw = genStrokes(config, )

image_name = os.getcwd() + '/image.png'
alpha_name = os.getcwd() + '/alpha.png'
image = cv2.imread(image_name)
alpha = cv2.imread(alpha_name)
image=cv2.resize(image,(400,400))
alpha=cv2.resize(alpha,(400,400))
alpha = alpha[:,:,0] 

regions_list = []
for x in range(-10, 10, 1):
    for y in range(-10, 10, 1):
        regions_list.append([x*1.0/10,y*1.0/10])
regions = np.array(regions_list)
######################################################################
def matteSolver(time_step,counter, image_name):
    im_sz = 400
    im_ch = 3
    output_ch = 1
    cnt = counter
    strokemapname = './strokemaps/stroke_' + str(cnt) +'.png'
    systemCall = './SharedMatting -i ./image.png -t ./strokemaps/stroke_' + str(cnt) +'.png -a'
    alphaname = './results/' + 'matte_' +str(time_step)+'_' + str(cnt) + '.png'
    while(not os.path.exists(image_name) or not os.path.exists(strokemapname)):
        print ('matteSolver: %d'%cnt)
        time.sleep(0.1)
    os.system(systemCall)
    matte = cv2.imread('shared_matting_alpha.bmp')
    cv2.imwrite(alphaname, matte)
    matte = cv2.imread(alphaname, 0)
    matte = matte.reshape((1, im_sz, im_sz, 1))
    return matte
######################################################################
def min_max_regions(time_step,selected_strokemap=None):
        mse_1_arr = []
        cnt = 0
        while(cnt < 10*10):
            loc = regions[cnt]
            cnt += 1
            l_1_j = loc
            strokemap_1_j = draw.drawstroke(image.copy(), l_1_j, alpha, cnt)
            strokemap_1_j = np.uint8(strokemap_1_j)
            strokename = './strokemaps/' + 'stroke_'+ str(cnt) + '.png'
            cv2.imwrite(strokename, strokemap_1_j)
        cnt1 = 0
        while(cnt1 < 10*10):
            cnt1 += 1
            matteSolver(time_step,cnt1, image_name)
        cnt2 = 0
        while(cnt2 < 10*10):
            cnt2 += 1
            alphaName = './results/' + 'matte_' +str(time_step)+'_' + str(cnt2) + '.png'
            alphaPNG = cv2.imread(alphaName)
            alpha_1_j = alphaPNG[:,:,0]
            alpha_1_j = alpha_1_j.astype(np.float)
            alpha_1_j = np.expand_dims(alpha_1_j, 0)
            mse_1_j = np.sqrt(np.sum(np.square(alpha - alpha_1_j)))
            mse_1_arr.append(mse_1_j)
        min_index = np.argmin(mse_1_arr)
        min_mse_1 = mse_1_arr[min_index]
        return min_mse_1, min_index
######################################################################
def min_max_regions_i(time_step,selected_strokemap):
        mse_i_arr = []
        cnt = 0
        while(cnt < 10*10):
            loc = regions[cnt]
            cnt += 1
            l_i_j = loc
            strokemap_i_j = draw.drawstroke(selected_strokemap.copy(), l_i_j, alpha, cnt)
            strokemap_i_j = np.uint8(strokemap_i_j)
            strokename = './strokemaps/' + 'stroke_' + str(cnt) + '.png'
            cv2.imwrite(strokename, strokemap_i_j)
        cnt1 = 0
        while(cnt1 < 10*10):
            cnt1 += 1
            matteSolver(time_step,cnt1, image_name)
        cnt2 = 0
        while(cnt2 < 10*10):
            cnt2 += 1
            alphaName = './results/' + 'matte_' +str(time_step)+'_' + str(cnt2) + '.png'
            alphaPNG = cv2.imread(alphaName)
            alpha_i_j = alphaPNG[:,:,0]
            alpha_i_j = alpha_i_j.astype(np.float)
            alpha_i_j = np.expand_dims(alpha_i_j, 0)
            mse_i_j = np.sqrt(np.sum(np.square(alpha - alpha_i_j)))
            mse_i_arr.append(mse_i_j)
        min_index = np.argmin(mse_i_arr)
        min_mse_i = mse_i_arr[min_index]
        return min_mse_i, min_index
######################################################################
def optimal_seq(index_arr):
    max_length = 30
    idx = 0
    while(idx < max_length):
        idx += 1
        if idx == 1:
            min_mse_1, min_index = min_max_regions(time_step=idx)
            selected_index = min_index
            print('selected index: %d' % selected_index)
            selected_strokename = './strokemaps/' + 'stroke_' + str(selected_index+1) + '.png'
            selected_strokemap = cv2.imread(selected_strokename)

        else:
            min_mse_i, min_index = min_max_regions_i(time_step=idx,selected_strokemap=selected_strokemap)
            selected_index = min_index
            index_arr.append(selected_index)
            print('selected index: %d' % selected_index)
            selected_strokename = './strokemaps/' + 'stroke_' + str(selected_index+1) + '.png'
            selected_strokemap = cv2.imread(selected_strokename)

if __name__ == '__main__':
    index_arr=[]
    optimal_seq(index_arr)
    print 'result is :', index_arr
    
